package com.zzh.partnersys.ai.chat;

import cn.hutool.core.collection.CollUtil;
import com.zzh.partnersys.ai.advisor.MyLoggerAdvisor;
import com.zzh.partnersys.ai.advisor.TokenUsedAdvisor;
import com.zzh.partnersys.ai.entity.AiAssistantDO;
import com.zzh.partnersys.ai.service.AiAssistantService;
import com.zzh.partnersys.ai.util.ThreadLocalContext;
import com.zzh.partnersys.common.enums.StatusEnum;
import com.zzh.partnersys.common.util.TokenUtils;
import jakarta.annotation.Resource;
import jakarta.servlet.http.HttpServletRequest;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
// import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor; // 暂时注释，如果恢复 QuestionAnswerAdvisor 请取消注释
// import org.springframework.ai.vectorstore.SearchRequest; // 暂时注释，如果恢复 QuestionAnswerAdvisor 请取消注释
import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.stream.Collectors;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;
import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY;

/**
 * @author: zzh
 * @date: 2025/05/03 23:09:43
 * @version: 1.0
 */
@Component
@Slf4j
public class AiChatApp {

    @Resource
    private ChatClient chatClient;

    @Resource
    private AiAssistantService aiAssistantService;

    @Resource
    private VectorStore elasticSearchVectorStore;

    private static final String DEFAULT_SYSTEM_PROMPT = "你是一个%s的AI助手,专注于专业质询服务,你要扮的角色是%s";

    /**
     * 新增流式接口方法
     * @param message  用户输入的消息
     * @param chatId  会话ID
     * @return Flux<String> 聊天机器人回复的消息流
     */
    public Flux<String> doChatWithStream(String message, String chatId, Long assistantId, HttpServletRequest request) {
        //获取当前会话的角色，根据角色调用不同的聊天机器人
        if (assistantId == null){
            return Flux.just("助手不存在");
        }
        AiAssistantDO aiAssistantDO = aiAssistantService.getAssistantById(assistantId);
        if (aiAssistantDO == null){
            return Flux.just("助手不存在");
        }
        if (StatusEnum.ENABLED.getCode() != aiAssistantDO.getStatus()){
            return Flux.just("助手暂未启用");
        }
        String assistantName = aiAssistantDO.getAssistantName();
        String roleDesc = aiAssistantDO.getRoleDesc();

        // 查询相似问题的文档数据（RAG检索）
        // 根据助手ID过滤文档，只检索该助手关联的文档
        SearchRequest searchRequest = SearchRequest.builder()
                .query(message)
                .topK(8)
                .similarityThreshold(0.7)
                .build();
        ThreadLocalContext.setContext(TokenUtils.getCurrentUserIdOrThrow(request),assistantId);
        //todo 其实可以直接不去用vectorStore的！，直接写一个Manager来实现add和search就好了！
        List<Document> documentList = elasticSearchVectorStore.similaritySearch(searchRequest);
        System.out.println("ES-documentList:"+documentList);
        ThreadLocalContext.clearContext();
        List<String> contentList = documentList.stream().map(Document::getText).toList();

        //会话记忆 Advisor获取
        // 注意：如果遇到模板语法错误，可能是 QuestionAnswerAdvisor 的默认模板有问题
        // 可以暂时注释掉 QuestionAnswerAdvisor 来测试是否是它导致的问题
        return chatClient
                .prompt()
                .system(String.format(DEFAULT_SYSTEM_PROMPT,assistantName,roleDesc))
                .user(CollUtil.isEmpty(contentList) ? message :
                        "基于以下上下文回答这个问题：" + message + ",如果从上下文中找不到答案，你可以自行尝试回答问题：" + String.join("\n", contentList))
                .advisors(spec -> spec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId)
                        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10))
                // 开启日志，便于观察效果
                .advisors(new MyLoggerAdvisor(),new TokenUsedAdvisor())
//                .tools(toolCallbackProvider)
                .stream()
                .content();
    }

}

